import sys
import os
import time
import tensorflow as tf
import numpy as np
from modelscope.models.audio.tts import SambertHifigan
from modelscope.outputs import OutputKeys
from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks

#model_dir = os.path.abspath("./pretrain_yangmi")
model_dir = os.path.abspath(sys.argv[1])

custom_infer_abs = {
    'voice_name':
    'F7',
    'am_ckpt':
    os.path.join(model_dir, 'tmp_am', 'ckpt'),
    'am_config':
    os.path.join(model_dir, 'tmp_am', 'config.yaml'),
    'voc_ckpt':
    os.path.join(model_dir, 'orig_model', 'basemodel_16k', 'hifigan', 'ckpt'),
    'voc_config':
    os.path.join(model_dir, 'orig_model', 'basemodel_16k', 'hifigan',
             'config.yaml'),
    'audio_config':
    os.path.join(model_dir, 'data', 'audio_config.yaml'),
    'se_file':
    os.path.join(model_dir, 'data', 'se', 'se.npy')
}
kwargs = {'custom_ckpt': custom_infer_abs}

model_id = SambertHifigan(os.path.join(model_dir, "orig_model"), **kwargs)

inference = pipeline(task=Tasks.text_to_speech, model=model_id, device='cpu')
srctxt = sys.argv[2]
dstpath = sys.argv[3]
with open(srctxt, 'r') as f:
   alltxt = f.readlines() 
for txtline in alltxt:
    txtline=txtline.strip().split()
    start = time.time()
    output = inference(input=txtline[1])
    print('spend time {}'.format(time.time()-start))
    wav = output[OutputKeys.OUTPUT_WAV]
    with open(os.path.join(dstpath,txtline[0]+".wav"), 'wb') as f:
        f.write(wav)

for syms in model_id.voices['F7'].tmp_symsinput:
    outsyms = '{'
    syms = syms.strip().split(" ")
    tmpsym = ""
    for sym in syms:
        sym = sym.strip("{").strip("}").split("$")
        print(sym)
        tmpsym += '{'
        for ph in sym:
            tmpsym += '"' + ph + '"'
            tmpsym += ','
        tmpsym = tmpsym[:-1]
        tmpsym += '},'
    outsyms += (tmpsym[:-1] + '}')
    print(outsyms)


models_path = sys.argv[4]
enc_model = tf.lite.Interpreter(os.path.join(models_path,'am/txtenc.tflite'))
emot_model = tf.lite.Interpreter(os.path.join(models_path,'am/emot.tflite'))
varadp_model1 = tf.lite.Interpreter(os.path.join(models_path,'am/varadp.tflite'))
duration_model = tf.lite.Interpreter(os.path.join(models_path,'am/duration.tflite'))
LR_model = tf.lite.Interpreter(os.path.join(models_path,'am_ext/LR.tflite'))
dursin_model = tf.lite.Interpreter(os.path.join(models_path,'am_ext/dursin.tflite'))
mem_model = tf.lite.Interpreter(os.path.join(models_path,'am_ext/mem.tflite'))
hdec_model = tf.lite.Interpreter(os.path.join(models_path,'am/attenH.tflite'))
getmask_model = tf.lite.Interpreter(os.path.join(models_path,'am_ext/getmask.tflite'))
dec1_model = tf.lite.Interpreter(os.path.join(models_path,'am/meldec1.tflite'))
dec2_model = tf.lite.Interpreter(os.path.join(models_path,'am/meldec2.tflite'))
#post_model = tf.lite.Interpreter(os.path.join(models_path,'am/post.tflite'))
dnorm_model = tf.lite.Interpreter(os.path.join(models_path,'am_ext/dnorm.tflite'))
#enc_model = tf.lite.Interpreter(os.path.join(models_path,'txtenc.tflite'))
#emot_model = tf.lite.Interpreter(os.path.join(models_path,'emot.tflite'))
#varadp_model1 = tf.lite.Interpreter(os.path.join(models_path,'varadp.tflite'))
#duration_model = tf.lite.Interpreter(os.path.join(models_path,'duration.tflite'))
#LR_model = tf.lite.Interpreter(os.path.join(models_path,'amext/LR.tflite'))
#dursin_model = tf.lite.Interpreter(os.path.join(models_path,'amext/dursin.tflite'))
#mem_model = tf.lite.Interpreter(os.path.join(models_path,'amext/mem.tflite'))
#hdec_model = tf.lite.Interpreter(os.path.join(models_path,'attenH.tflite'))
#getmask_model = tf.lite.Interpreter(os.path.join(models_path,'amext/getmask.tflite'))
#dec1_model = tf.lite.Interpreter(os.path.join(models_path,'meldec1.tflite'))
#dec2_model = tf.lite.Interpreter(os.path.join(models_path,'meldec2.tflite'))
#post_model = tf.lite.Interpreter(os.path.join(models_path,'post.tflite'))
#dnorm_model = tf.lite.Interpreter(os.path.join(models_path,'amext/dnorm.tflite'))

def tf_tts_infer(am_input):
    inputs_ling = am_input[0].cpu().numpy()
    sig = enc_model.get_signature_runner('serving_default')
    sig_out = sig(text_in=tf.convert_to_tensor(inputs_ling,dtype=tf.int32))
    enc_out = sig_out['text_fft_encoder']
    #print('enc out {} {}'.format(enc_out.shape,enc_out))
    return enc_out

def tf_variance_adaptor(am_input, txt_hid):
    inputs_emotion = am_input[1].cpu().numpy()
    inputs_emotion = tf.convert_to_tensor(inputs_emotion,dtype=tf.int32)
    sig = emot_model.get_signature_runner('serving_default')
    sig_out = sig(args_0 = inputs_emotion)
    emo_hid = sig_out['emot_emd']

    spk_hid = tf.convert_to_tensor(am_input[2].cpu().numpy())
    print('emohid {} {}'.format(emo_hid.shape,emo_hid[:,:,:]))
    print('spkhid {} {}'.format(emo_hid.shape,spk_hid[:,:,:]))

    print('txt_hid {} {}'.format(txt_hid.shape,txt_hid[:,:,:]))

    sig = varadp_model1.get_signature_runner('serving_default')
    sig_out = sig(emo_in=emo_hid,spk_in=spk_hid,txt_in=txt_hid)
    txt_emb_aug = sig_out['variance_adaptor1']
    input_pred_con = sig_out['variance_adaptor1_1']
    t_emd = sig_out['variance_adaptor1_2']
    p_emd = sig_out['variance_adaptor1_3']
    e_emd = sig_out['variance_adaptor1_4']
    print('txt_emb_aug {} {}'.format(txt_emb_aug.shape,txt_emb_aug))
    print('input_pred_con {} {}'.format(input_pred_con.shape,input_pred_con))
    print('temd {} {}'.format(t_emd.shape,t_emd))
    print('pemd {} {}'.format(p_emd.shape,p_emd))
    print('eemd {} {}'.format(e_emd.shape,e_emd))

    sig = duration_model.get_signature_runner('serving_default')
    cond_len = tf.shape(input_pred_con).numpy()[1]
    pre_x = tf.zeros([1,1,1])
    pre_h0 = tf.zeros([1,128])
    pre_c0 = tf.zeros([1,128])
    pre_h1 = tf.zeros([1,128])
    pre_c1 = tf.zeros([1,128])
    dura_output = []
    for il in range(cond_len):
        sig_out = sig(cond = input_pred_con[:,il:il+1,:],p_c0=pre_c0,p_c1=pre_c1,p_h0=pre_h0,p_h1=pre_h1,txt_in=pre_x)
        out_x=sig_out['var_rnn_ar_predictor']
        h0=sig_out['var_rnn_ar_predictor_1']
        c0=sig_out['var_rnn_ar_predictor_2']
        h1=sig_out['var_rnn_ar_predictor_3']
        c1=sig_out['var_rnn_ar_predictor_4']
        out_logx=sig_out['var_rnn_ar_predictor_5']
        dura_output.append(out_logx)
        pre_x = out_x
        pre_h0 = h0
        pre_c0 = c0
        pre_h1 = h1
        pre_c1 = c1
    duration_predictions = tf.concat(dura_output,axis=1)

    sig = LR_model.get_signature_runner('serving_default')
    sig_out = sig(lr_inputs=txt_emb_aug,duration_pred=duration_predictions)
    LR_txt_out = sig_out['length_regulator']
    LR_txt_len = sig_out['length_regulator_1']

    sig_out = sig(lr_inputs=emo_hid,duration_pred=duration_predictions)
    LR_emo_out = sig_out['length_regulator']
    sig_out = sig(lr_inputs=spk_hid,duration_pred=duration_predictions)
    LR_spk_out = sig_out['length_regulator']

    sig = dursin_model.get_signature_runner('serving_default')
    sig_out = sig(duration_pred=duration_predictions,lr_txt_in=LR_txt_out)
    LR_txt_out = sig_out['dur_sinusoidal_position_encoder']


    sig = mem_model.get_signature_runner('serving_default')
    sig_out = sig(txt_in=LR_txt_out,spk_in=LR_spk_out,emo_in=LR_emo_out)
    memout = sig_out['txtmem']
    print('memout {} {}'.format(memout.shape,memout))
    return memout,duration_predictions,LR_txt_len

def tf_mel_dec(memory,dur_pred):
    sig = hdec_model.get_signature_runner('serving_default')
    start = time.time()
    sig_out = sig(mem_in = memory)
    hdec_out1 = sig_out['hybrid_attention_decoder_h']
    hdec_out2 = sig_out['hybrid_attention_decoder_h_1']

    dur_maxlen = int(tf.math.reduce_max(dur_pred)/3+0.5)
    mem_len = tf.shape(memory)[1]
    sig = getmask_model.get_signature_runner('serving_default')
    sig_out = sig(max_in=tf.convert_to_tensor([[mem_len]]),x_in=tf.convert_to_tensor([[dur_maxlen]]),h_in=tf.convert_to_tensor([[dur_maxlen]]))
    x_mask = sig_out['getmask']
    h_mask = sig_out['getmask_1']

    sig_dec1 = dec1_model.get_signature_runner('serving_default')
    sig_dec2 = dec2_model.get_signature_runner('serving_default')
    mem_len = tf.shape(memory)[1].numpy()
    dec_inputs = tf.zeros([1,1,82])
    p_xk = tf.zeros([12,8,1,16])
    p_xv = tf.zeros([12,8,1,16])
    dec_outlist = []
    for step in range(mem_len):
        memtmp = memory[:,step:step+1,:]
        #print('step {} memin {} {}'.format(step,memtmp[:,:,:5],memtmp[:,:,-5:]))
        #print('step {} dec_inputs {} {}'.format(step,dec_inputs[:,:,:5],dec_inputs[:,:,-5:]))
        sig_dec1_out = sig_dec1(mel_in=dec_inputs,mem_in=memory[:,step:step+1,:])
        dec1_output = sig_dec1_out['hybrid_attention_decoder_1']
        #print('step {} dec1 {} {}'.format(step,dec1_output[:,:,:5],dec1_output[:,:,-5:]))
        #print('step {} xqin  {} {}'.format(step,dec1_output[:,:,:5],dec1_output[:,:,-5:]))
        #print('step {} xkin  {} {}'.format(step,p_xk[2,0,step,:5],p_xk[2,0,step,-5:]))
        #print('step {} xvin  {} {}'.format(step,p_xv[2,0,step,:5],p_xv[2,0,step,-5:]))
        #print('step {} hkin  {} {}'.format(step,hdec_out1[2,0,step,:5],hdec_out1[2,0,step,-5:]))
        #print('step {} hvin  {} {}'.format(step,hdec_out2[2,0,step,:5],hdec_out2[2,0,step,-5:]))
        #print('step {} xmask  {}'.format(step,x_mask[:,step:step+1,:10]))
        #print('step {} hmask  {}'.format(step,h_mask[:,step:step+1,:10]))
        sig_dec2_out = sig_dec2(x_q_in=dec1_output,x_k_in=p_xk,x_v_in=p_xv,h_k_in=hdec_out1,h_v_in=hdec_out2,x_mask_in=x_mask[:,step:step+1,:step+1],h_mask_in=h_mask[:,step:step+1,:])
        dec2_output=sig_dec2_out['hybrid_attention_decoder_2']
        x_k=sig_dec2_out['hybrid_attention_decoder_2_1']
        x_v=sig_dec2_out['hybrid_attention_decoder_2_2']
        pnc_output=sig_dec2_out['hybrid_attention_decoder_2_3']
        norm_output=sig_dec2_out['hybrid_attention_decoder_2_4']
        lin_output=sig_dec2_out['hybrid_attention_decoder_2_5']
        declayer_output=sig_dec2_out['hybrid_attention_decoder_2_6']
        dec_inputs = dec2_output[:,:,-82:]
        #print('step {} dec2 out {} {}'.format(step,dec2_output[:,:,:5],dec2_output[:,:,-5:]))
        #print('step {} declayer out {} {}'.format(step,declayer_output[1,:,:5],declayer_output[1,:,-5:]))
        #print('step {} declayer out {} {}'.format(step,declayer_output[2,:,:5],declayer_output[2,:,-5:]))
        #print('step {} pnc out {} {}'.format(step,pnc_output[0,:,:5],pnc_output[0,:,-5:]))
        #print('step {} norm out {} {}'.format(step,norm_output[0,:,:5],norm_output[0,:,-5:]))
        #print('step {} lin out {} {}'.format(step,lin_output[0,:,:5],lin_output[0,:,-5:]))
        p_xk = x_k
        p_xv = x_v
        dec_outlist.append(dec2_output)
        #print('dec output {} {}'.format(dec2_output[:,:,:5],dec2_output[:,:,-5:]))    
    decout = tf.concat(dec_outlist,axis=1)
    #print('dec output {}'.format(decout))    
    return decout

def tf_mel_post(mel_dec_out,valid_len):
    mel_dec_out = tf.reshape(mel_dec_out,(1,-1,82))
    #for i in range(mel_dec_out.shape[1]):
    #    print('postin {} {} {}'.format(i, mel_dec_out[:,i,:5],mel_dec_out[:,i,-5:]))
    valid_len = valid_len[0]
    if valid_len < tf.shape(mel_dec_out)[1]:
        mel_dec_out = tf.concat([mel_dec_out[:,:valid_len,:],tf.zeros([1,tf.shape(mel_dec_out)[1]-valid_len,82])],axis=1)
        postmask = tf.concat([tf.ones([1,valid_len,1],dtype=tf.bool),tf.zeros([1,tf.shape(mel_dec_out)[1]-valid_len,1],dtype=bool)],axis=1)
    else:
        postmask = tf.ones([1,valid_len,1],dtype=tf.bool)

    post_model = tf.lite.Interpreter(os.path.join(models_path,'am/post0.tflite'))
    sig = post_model.get_signature_runner('serving_default')
    sig_out = sig(post_in=mel_dec_out,mask=postmask,lpadding=tf.zeros([1,37,256]),rpadding=tf.zeros([1,3,256]))
    post_output = sig_out['post_net0']
    post_model = tf.lite.Interpreter(os.path.join(models_path,'am/post1.tflite'))
    sig = post_model.get_signature_runner('serving_default')
    sig_out = sig(post_in=post_output,mask=postmask,lpadding=tf.zeros([1,37,256]),rpadding=tf.zeros([1,3,256]))
    post_output = sig_out['post_net1']
    post_model = tf.lite.Interpreter(os.path.join(models_path,'am/post2.tflite'))
    sig = post_model.get_signature_runner('serving_default')
    sig_out = sig(post_in=post_output,mask=postmask,lpadding=tf.zeros([1,37,256]),rpadding=tf.zeros([1,3,256]))
    post_output = sig_out['post_net2']
    post_model = tf.lite.Interpreter(os.path.join(models_path,'am/post3.tflite'))
    sig = post_model.get_signature_runner('serving_default')
    sig_out = sig(post_in=post_output,mask=postmask,lpadding=tf.zeros([1,37,256]),rpadding=tf.zeros([1,3,256]))
    post_output = sig_out['post_net3']

    post_model = tf.lite.Interpreter(os.path.join(models_path,'am/post.tflite'))
    sig = post_model.get_signature_runner('serving_default')
    sig_out = sig(post_in=post_output,p_h=tf.zeros([1,128]),p_c=tf.zeros([1,128]),post_res=mel_dec_out)
    post_output = sig_out['post_net']
    post_output = tf.where(postmask,post_output,tf.zeros(1))
    #for i in range(valid_len):
    #    print('postout {} {} {}'.format(i, post_output[:,i,:5],post_output[:,i,-5:]))
    return post_output[:,:valid_len,:]

def denorm_f0(mel):
    sig = dnorm_model.get_signature_runner('serving_default')
    sig_out = sig(norm_in=mel)
    res_mel = sig_out['dnorm']

    #print('dnorm {} {}'.format(res_mel[:,:,:5],res_mel[:,:,-5:]))
    for i in range(res_mel.shape[1]):
        print('dnorm {} {} {}'.format(i, res_mel[:,i,:5],res_mel[:,i,-5:]))

    from tts_voc_tflite_mul import voc_infer
    #from tts_voc_loadmul import voc_infer
    audio = voc_infer(res_mel[0,:,:])
    return audio 

outpcms = []
for kid,aminput in enumerate(model_id.voices['F7'].tmp_aminput[:1]):
    txt_embed = tf_tts_infer(aminput)
    tf_var_out,dur_prdiction,LR_length_rounded = tf_variance_adaptor(aminput,txt_embed)
    mel_output = tf_mel_dec(tf_var_out,dur_prdiction)
    post_mel = tf_mel_post(mel_output,LR_length_rounded)
    pcm = denorm_f0(post_mel)
    outpcms.append(pcm)
outpcms = np.concatenate(outpcms,axis=0)
import wave
with wave.open(dstpath+"/tflitedebug.wav", 'wb') as wf:
    wf.setnchannels(1)
    wf.setsampwidth(2)
    wf.setframerate(16000)
    wf.writeframes(outpcms.data)
